import signal
import requests
from aip import AipSpeech
import speech_recognition as sr
from pydub import AudioSegment
from pydub.playback import play
import os
import tempfile
import io
import numpy as np
import asyncio
import threading
import websockets
import json
import sys
import argparse

from TTS import TTS_By_BaiduAPI

# Import the OpenCC library for text conversion
from opencc import OpenCC
cc = OpenCC('t2s')

# Global variables for WebSocket management
connected_clients = set()
clients_lock = threading.Lock()
loop = None
stop_event = threading.Event()
# ASR control variables
asr_active = False
asr_control_event = threading.Event()
# Server URL
SERVER_URL = "http://123.60.35.244:8000/transcribe"
WEBSOCKET_URL = "192.168.250.231"

async def run_server():
    # Start the WebSocket server and keep it running
    async with websockets.serve(
        handler,
        WEBSOCKET_URL, 
        8765,
        # 关键参数：开启跨域支持，允许所有来源
        origins=None
    ):
        while not stop_event.is_set():
            await asyncio.sleep(2)

# WebSocket handler to manage client connections
async def handler(websocket):
    """Handle WebSocket connections and process incoming messages."""
    global asr_active
    
    with clients_lock:
        print("Client connected")
        connected_clients.add(websocket)    
    try:
        # Process incoming messages
        async for message in websocket:
            try:
                data = json.loads(message)
                if "voice" in data:
                    if data["voice"] == "on":
                        asr_active = True
                        asr_control_event.set()  # Signal ASR to start
                        print("ASR activated")
                        await websocket.send(json.dumps({"status": "ASR activated"}))
                    elif data["voice"] == "off":
                        asr_active = False
                        print("ASR deactivated")
                        await websocket.send(json.dumps({"status": "ASR deactivated"}))
            except json.JSONDecodeError:
                print(f"Received invalid JSON: {message}")
                
        await websocket.wait_closed()
    finally:
        with clients_lock:
            connected_clients.remove(websocket)

# Signal handler to initiate graceful shutdown
def signal_handler(sig, frame):
    print("Shutting down gracefully...")
    stop_event.set()  # Signal the server to stop
    asr_active = False
    asr_control_event.set()  # Wake up ASR thread if it's waiting
    websocket_thread.join()  # Wait for the WebSocket thread to finish
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)

# Start the WebSocket server in a separate thread
def start_websocket_server():
    global loop
    # Create a new event loop for this thread
    loop = asyncio.new_event_loop()
    # Set it as the current event loop for the thread
    asyncio.set_event_loop(loop)
    # Run the server coroutine, which will keep the loop running
    loop.run_until_complete(run_server())

# Send a message to all connected WebSocket clients
def send_message(message):
    """Send a JSON message to all connected clients."""
    with clients_lock:
        for client in connected_clients:
            asyncio.run_coroutine_threadsafe(client.send(message), loop)

def process_misspellings(text):
    print("原始文本: " + text)
    if "爱国荣耀" in text:
        text = text.replace("爱国荣耀", "爱国荣校")
    if "思源" in text and "爱国" in text:
        text = "饮水思源爱国荣校"
    if "诗源" in text:
        text = text.replace("诗源", "思源")
    if "资源" in text:
        text = text.replace("资源", "思源")
    if "寺院" in text:
        text = text.replace("寺院", "思源")
    if "词源" in text:
        text = text.replace("词源", "思源")
    if "银水思源" in text:
        text = text.replace("银水思源", "饮水思源")
    if "银水私源" in text:
        text = text.replace("银水私源", "饮水思源")
    return text

def process_misspellings_original(text):
    # print("原始文本: " + text)
    # check if the length bigger than 18 characters
    if len(text) > 18:
        text = "我要写入饮水思源"
    if "爱国荣耀" in text:
        text = text.replace("爱国荣耀", "爱国荣校")
    if "思源" in text and "爱国" in text:
        text = "我要写入饮水思源爱国荣校"
    if "诗源" in text:
        text = text.replace("诗源", "思源")
    if "资源" in text:
        text = text.replace("资源", "思源")
    if "寺院" in text:
        text = text.replace("寺院", "思源")
    if "词源" in text:
        text = text.replace("词源", "思源")
    if "银水思源" in text:
        text = text.replace("银水思源", "饮水思源")
    if "银水私源" in text:
        text = text.replace("银水私源", "饮水思源")
    return text

def process_misspellings_duqu(text):
    # print("原始文本: " + text)
    # check if the length bigger than 18 characters
    if len(text) > 18:
        text = "我要读取数据"
    if "爱国荣耀" in text:
        text = text.replace("爱国荣耀", "爱国荣校")
    if "思源" in text and "爱国" in text:
        text = "我要读取饮水思源爱国荣校"
    if "诗源" in text:
        text = text.replace("诗源", "思源")
    if "资源" in text:
        text = text.replace("资源", "思源")
    if "寺院" in text:
        text = text.replace("寺院", "思源")
    if "词源" in text:
        text = text.replace("词源", "思源")
    if "银水思源" in text:
        text = text.replace("银水思源", "饮水思源")
    if "银水私源" in text:
        text = text.replace("银水私源", "饮水思源")
    return text

def filter_audio(audio_data):
    # Apply band-pass filter (300 Hz to 3400 Hz)
    audio_segment = AudioSegment.from_wav(io.BytesIO(audio_data))
    filtered_audio = audio_segment.band_pass_filter(min_freq=300, max_freq=3400)
    filtered_audio_data = filtered_audio.export(format="wav").read()
    return filtered_audio_data

def task(text: str):
    """Process transcribed text and send text/answer pairs to WebSocket clients."""
    if text and "同学" in text:
        answer = "你好，我是思源同学"
        message = json.dumps({"text": text, "answer": answer})
        send_message(message)
        TTS_By_BaiduAPI("你好，我是思源同学")
        
    if text and "读取" in text:
        text = process_misspellings_duqu(text)
        text_to_write = text.replace("我要读取", "").replace("请读取","").strip()
        text_to_write = process_misspellings(text_to_write)
        answer = f"正在读取: {text_to_write}"
        message = json.dumps({"text": text, "answer": answer})
        send_message(message)
        TTS_By_BaiduAPI("思源同学正在为您读取: " + (text_to_write))

    if text and "雪若" in text:
        text.replace("雪若", "写入")
    if text and "写入" in text:
        text = process_misspellings_original(text)
        text_to_write = text.replace("我要写入", "").replace("请写入","").strip()

        text_to_write = process_misspellings(text_to_write)
        answer = f"正在写入: {text_to_write}"
        
        message = json.dumps({"text": text, "answer": answer})
        send_message(message)
        TTS_By_BaiduAPI("思源同学正在为您写入: " + (text_to_write))

def ASR_By_BaiduAPI():
    """Speech recognition using Baidu API with WebSocket integration."""
    # Baidu API credentials
    APP_ID = '*'
    API_KEY = '*'
    SECRET_KEY = '*'
    
    client = AipSpeech(APP_ID, API_KEY, SECRET_KEY)
    
    def get_text(wav_bytes):
        result = client.asr(wav_bytes, 'wav', 16000, {'dev_pid': 1537})
        try:
            text = result["result"][0]
        except Exception as e:
            print(e)
            text = ""
        return text
    
    r = sr.Recognizer()
    mic = sr.Microphone()
    
    while not stop_event.is_set():
        # Check if ASR is active
        if not asr_active:
            print("ASR is inactive. Waiting for activation...")
            # Wait until ASR is activated or program is shutting down
            asr_control_event.clear()
            asr_control_event.wait()
            continue
            
        print("请说话...")
        with mic as source:
            # r.adjust_for_ambient_noise(source, duration=0.5)
            r.energy_threshold = max(1000, r.energy_threshold)
            try:
                audio = r.listen(source, timeout=10, phrase_time_limit=30)
            except sr.WaitTimeoutError:
                print("10秒内未听到您讲话。")
                continue
        
        # Check if ASR is still active after listening
        if not asr_active:
            continue
            
        audio_data = audio.get_wav_data(convert_rate=16000)
        print("\n正在分析...")
        
        text = get_text(audio_data)
        print("\n识别结果: " + text)
        
        task(text)

def ASR_By_whisper():
    """Speech recognition using Whisper with WebSocket integration."""
    
    r = sr.Recognizer()
    mic = sr.Microphone()
    r.pause_threshold = 1.1
    
    while not stop_event.is_set():
        # Check if ASR is active
        if not asr_active:
            print("ASR is inactive. Waiting for activation...")
            # Wait until ASR is activated or program is shutting down
            asr_control_event.clear()
            asr_control_event.wait()
            continue
        
        with mic as source:
            r.adjust_for_ambient_noise(source, duration=0.4)
            print("请说话...")
            r.energy_threshold = max(1000, 1.1 * r.energy_threshold)
            try:
                audio = r.listen(source, timeout=10, phrase_time_limit=30)
            except sr.WaitTimeoutError:
                print("10秒内未听到您讲话。")
                continue
                
        # Check if ASR is still active after listening
        if not asr_active:
            continue
            
        audio_data = audio.get_wav_data(convert_rate=16000)
        print("\n正在识别中...")
        
        try:
            # Send audio data to the cloud server
            response = requests.post(
                SERVER_URL,
                files={"file": ("audio.wav", audio_data, "audio/wav")},
                timeout=30
            )
            if response.status_code == 200:
                result = response.json()
                if "transcription" in result:
                    text = result["transcription"]
                elif "error" in result:
                    print(f"服务器错误: {result['error']}")
                    text = ""
                else:
                    print("未知的服务器响应")
                    text = ""
            else:
                print(f"HTTP错误: {response.status_code}")
                text = ""
        except requests.exceptions.Timeout:
            print("请求超时")
            text = ""
        except requests.exceptions.RequestException as e:
            print(f"请求失败: {e}")
            text = ""
        
        print("识别结果: \n" + text)
        task(text)

def play_audio(audio):
    """Play an audio file."""
    try:
        print("正在播放合成音频...")
        sound = AudioSegment.from_file(audio)
        play(sound)
        print("播放结束")
    except Exception as e:
        print(f"音频播放失败：{str(e)}（请检查音频文件格式或播放库依赖）")

if __name__ == "__main__":
    # parse command line arguments
    # if --model = baidu then use baidu api
    # if --model = whisper then use whisper
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model", type=str, default="whisper", help="Choose ASR model (baidu or whisper)")
    parser.add_argument("-o", "--on", action='store_true', help="Turn on ASR immediately")
    args = parser.parse_args()
    asr_active = args.on
    # Start WebSocket server in a separate thread
    websocket_thread = threading.Thread(target=start_websocket_server)
    websocket_thread.start()

    if args.model == "baidu":
        func_to_run = ASR_By_BaiduAPI
    else:
        func_to_run = ASR_By_whisper
    try:
        func_to_run()
    except KeyboardInterrupt:
        websocket_thread.join()
    print("Program exited cleanly.")